#!/usr/bin/env python

import rospy
from sensor_msgs.msg import Image
from geometry_msgs.msg import Twist
from cv_bridge import CvBridge
import cv2
import torch
import sys
import os
# 获取当前脚本的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# 构建相对路径的绝对路径
relative_path = os.path.join(current_dir, 'yolov5')
# 添加到sys.path
sys.path.append(relative_path)

from yolov5.models.experimental import attempt_load
from yolov5.utils.general import non_max_suppression
from yolov5.models.common import DetectMultiBackend


def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):
    # 获取图像的原始尺寸
    shape = img.shape[:2]  # current shape [height, width]
    
    # 计算缩放比例
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
    
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better test mAP)
        r = min(r, 1.0)
    
    # 计算缩放后的尺寸
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, 32), np.mod(dh, 32)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
    
    dw /= 2  # divide padding into 2 sides
    dh /= 2
    
    if shape[::-1] != new_unpad:  # resize
        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    
    return img, ratio, (dw, dh)
class Yolov5Detector:
    def __init__(self):
        rospy.init_node('yolov5_detector', anonymous=True)
        
        # 加载YOLOv5模型
        model_path = current_dir + '/yolov5/yolov5s.pt'
        self.model = attempt_load(model_path)
        self.model.eval()
      
        # 初始化CvBridge
        self.bridge = CvBridge()
        
        # 订阅图像话题
        # self.image_sub = rospy.Subscriber('/usb_cam/image_raw', Image, self.image_callback)
       
        # self.image_sub = rospy.Subscriber( "/kinect2/sd/image_ir_rect", Image, self.image_callback)
        self.image_sub = rospy.Subscriber( "/kinect2/hd/image_color_rect", Image, self.image_callback)
        # 发布检测结果话题
        self.detection_pub = rospy.Publisher('/yolov5/detections', Image, queue_size=10)
        # 发布速度话题
        self.vel_pub = rospy.Publisher('/cmd_vel', Twist, queue_size=10)

    def image_callback(self, msg):
        # 将ROS图像消息转换为OpenCV图像
        cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
        # 获取图像的大小
        height, width, channels = cv_image.shape
        # 使用letterbox函数进行图像缩放和填充
        img, ratio, pad = letterbox(cv_image, (640, 640), auto=False)
        # 将图像转换为模型输入格式
        img = torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0) / 255.0
        ball_x,ball_y = None,None
        # 进行推理
        with torch.no_grad():
            pred = self.model(img)[0]
            pred = non_max_suppression(pred, 0.25, 0.45, max_det=1000)
        
        # 处理检测结果
        for det in pred:
            if det is not None and len(det):
                for *xyxy, conf, cls in reversed(det):
                    xyxy = [int(x) for x in xyxy]
                    xyxy[0] -= pad[0]  # x1
                    xyxy[1] -= pad[1]  # y1
                    xyxy[2] -= pad[0]  # x2
                    xyxy[3] -= pad[1]  # y2
                                       
                    # 分解ratio为单独的宽高比
                    ratio_w, ratio_h = ratio
                    
                    # 缩放回原始尺寸
                    xyxy[0] = int(xyxy[0] / ratio_w)
                    xyxy[1] = int(xyxy[1] / ratio_h)
                    xyxy[2] = int(xyxy[2] / ratio_w)
                    xyxy[3] = int(xyxy[3] / ratio_h)
                    ball_x = int((xyxy[0]+xyxy[2])/2)
                    ball_y = int((xyxy[1]+xyxy[3])/2)
                    label = f'{self.model.names[int(cls)]} {conf:.2f}'
                    cv2.rectangle(cv_image, (xyxy[0], xyxy[1]), (xyxy[2], xyxy[3]), (0, 255, 0), 2)                
                    cv2.putText(cv_image, label, (xyxy[0], xyxy[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
        # 发布检测结果图像
        detection_msg = self.bridge.cv2_to_imgmsg(cv_image, "bgr8")
        self.detection_pub.publish(detection_msg)
        cmd_msg = Twist()
        # TODO
        if(ball_x):        
            print(f"ball_x:{ball_x/width},ball_y:{ball_y/height}")
            cmd_msg = Twist()
            cmd_msg.linear.x = 0
            cmd_msg.angular.z = 0.9*(0.5-ball_x/width)
            if abs(ball_x/width - 0.5) < 0.1:
                print("go")
                # cmd_msg.linear.x = 1
                cmd_msg.linear.x = 0.9*(0.8-ball_y/height)   
        else:
            cmd_msg.linear.x = 0
            cmd_msg.angular.z = 0
        self.vel_pub.publish(cmd_msg) 


    def run(self):
        rospy.spin()

if __name__ == '__main__':
    detector = Yolov5Detector()
    detector.run()
